from mmcv.cnn.bricks import padding
from numpy.core.numeric import identity
import torch
from mmcv.ops.deform_conv import DeformConv2dPack, deform_conv2d
from mmcv.cnn import CONV_LAYERS, UPSAMPLE_LAYERS
import torch.nn as nn
import math
import torch.nn.functional as F

@UPSAMPLE_LAYERS.register_module("deform_weighted_kernel_v2")
class DeformUpsampleBlock(DeformConv2dPack):
    """
    Add Rep structure to DCN module, Block 意味着将batchnorm层融入该模块 
    """
    def __init__(self, scale_factor=2, aap_size=7, *args, **kwargs):
        super().__init__(*args, **kwargs)
        assert self.out_channels == self.in_channels, "out_channel must match in_channels"
        # self.groups = self.in_channels   #分组卷积
        self.scale = scale_factor # 上采样尺度
        self.aap_size = aap_size
        # self.temperature = 10
        del self.weight
        # only weight, no bias
        self.conv_weight = nn.ModuleList(
            [nn.AdaptiveAvgPool2d(self.aap_size),
            nn.Conv2d(
                in_channels=self.in_channels, 
                out_channels=self.kernel_size[0]**2, 
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False
            ),
            nn.Conv2d(
                in_channels=self.aap_size**2, 
                out_channels=self.scale**2, 
                kernel_size=1,
                stride=1,
                padding=0,
                bias=False
            )]
        )
        self.PS = torch.nn.PixelShuffle(self.scale)


    def forward(self, x):
        # 正确版本
        offset = self.conv_offset(x) # [batch, 18, rows, cols]
        # 对weight使用softmax时会导致梯度为None，需要小心，注意self.weight于self.weight.data的差别

        [B, C, H, W] = x.size()
        # weight 变换过程 [2,256,H,W] -> [2,9,7,7] -> [2,49,3,3] -> [2,4,3,3]"
        self.weight0 = self.conv_weight[0](x)
        self.weight1 = self.conv_weight[1](self.weight0)
        self.weight1_permute = self.weight1.reshape(B, self.kernel_size[0], self.kernel_size[1], self.aap_size**2).permute(0,3,1,2).contiguous()
        self.weight = self.conv_weight[2](self.weight1_permute)
        
        weight_tmp_ = self.weight.reshape(-1, self.scale**2, 1,self.kernel_size[0]*self.kernel_size[1])
        weight_tmp = F.softmax(weight_tmp_, dim=-1)
        weight_tmp = weight_tmp.reshape(-1, self.scale**2, 1, self.kernel_size[0], self.kernel_size[1])
        weight_repeat = weight_tmp.repeat(1, self.out_channels, 1, 1, 1)
        weight_last = weight_repeat.reshape(-1, 1, self.kernel_size[0], self.kernel_size[1])
        x = x.reshape(-1, B*C, H, W)
        offset = offset.reshape(-1, B*self.kernel_size[0] * self.kernel_size[1] * 2, H, W)
        output = deform_conv2d(x, offset, weight_last, self.stride, self.padding,
                            self.dilation, self.groups*B, self.deform_groups*B)
        output = output.reshape(B, -1 ,H, W)
        output = self.PS(output)
        return output



if __name__ == "__main__":    
    model = DeformUpsampleBlock(in_channels=256,out_channels=256,kernel_size=(3,3), padding=1, groups=256, aap_size=9).cuda()
    input = torch.randn(2,256,4,4).cuda()
    out1 = model.forward(input)
    print(out1.shape)
    out2 = model.forward2(input)
    # out3 = model.forward3(input)
    # print((out2 == out1).all())
    # print((out2 == out3).all())
